#include <chrono>
#include <condition_variable>
#include <functional>
#include <iostream>
#include <mutex>
#include <queue>
#include <thread>
#include <vector>

template <class T>
class ThreadPool {
 public:
  ThreadPool(int thread_num) : thread_num_(thread_num) {
    start_ = true;
    for (int i = 0; i < thread_num; ++i) {
      std::thread new_thread(&ThreadPool::DoSomething, this);
      pool_.push_back(std::move(new_thread));
      std::cout << "add new thread! " << std::endl;
    }
  }
  ~ThreadPool(){
      ClearTask();
  }
  void AddTask(T task) {
      {
          std::unique_lock<std::mutex> lk(mutex_);
          task_.push(task);
      }
      cond_var_.notify_one();
  }
  void Stop() {
    start_ = false;
    cond_var_.notify_all();
    for (auto& thread : pool_) {
      thread.join();
    }
  }
  void ClearTask(){
      while(!task_.empty()){
          task_.pop();
      }
  }
  void DoSomething() {
    while(true){
        std::unique_lock<std::mutex> lk(mutex_);
        while (task_.empty() && start_) {
          std::cout << "worker thread waiting, thread id: "
                    << std::this_thread::get_id() << std::endl;
          cond_var_.wait(lk);
        }
        if(!start_){
            break;
        }
        auto task = task_.front();
        task_.pop();

        lk.unlock();
        std::cout << "worker thread get a task! thread id: "
                  << std::this_thread::get_id() << std::endl;
        std::invoke(task.first, task.second);

    }

    std::cout << "worker thread exit! thread id: " << std::this_thread::get_id()
              << std::endl;
  }

 private:
  bool start_;
  int thread_num_;
  std::vector<std::thread> pool_;
  std::queue<T> task_;
  std::mutex mutex_;
  std::condition_variable cond_var_;
};

void UserTask(int taskNum) {
  std::cout << " this is user task, task number :" << taskNum << std::endl;
  std::this_thread::sleep_for(std::chrono::milliseconds(10));
}

using func_type = std::function<void(int)>;
using task_type = std::pair<func_type, int>;

int main() {
  static int taskSum = 0;
  ThreadPool<task_type> thread_pool(4);
  for (int taskSum = 0; taskSum < 1000; ++taskSum) {
    thread_pool.AddTask({func_type(UserTask), taskSum});
  }
  std::this_thread::sleep_for(std::chrono::seconds(5));
  thread_pool.Stop();
  return 0;
}
